import torch.nn as nn
import torch.nn.functional as F


def create_cnn_with_bn_leaky(in_channels, out_channels, return_module=False, **kwargs):
    modules = [
        nn.Conv2d(in_channels, out_channels, bias=False, **kwargs),
        nn.BatchNorm2d(out_channels),
        nn.LeakyReLU(0.1, inplace=True),
    ]
    if return_module:
        return nn.Sequential(*modules)
    else:
        return modules


class GlobalAvgPool2d(nn.Module):
    def _init__(self):
        super(GlobalAvgPool2d, self).__init__()

    def forward(self, x):
        N, C, H, W = x.data.size()
        x = F.avg_pool2d(x, (H, W))
        x = x.view(N, C)
        return x


class Darknet19(nn.Module):
    architecture = {
        # fmt: off
        "layer0": [     (  32, 3, 1, 1)],
        "layer1": ["M", (  64, 3, 1, 1)],
        "layer2": ["M", ( 128, 3, 1, 1), ( 64, 1, 1, 0), ( 128, 3, 1, 1)],
        "layer3": ["M", ( 256, 3, 1, 1), (128, 1, 1, 0), ( 256, 3, 1, 1)],
        "layer4": ["M", ( 512, 3, 1, 1), (256, 1, 1, 0), ( 512, 3, 1, 1), (256, 1, 1, 0), ( 512, 3 ,1, 1)],
        "layer5": ["M", (1024, 3, 1, 1), (512, 1, 1, 0), (1024, 3, 1, 1), (512, 1, 1, 0), (1024, 3 ,1, 1)]
        # fmt: on
    }

    def __init__(self, in_channels=3, num_classes=1000):
        super(Darknet19, self).__init__()

        # fmt: off
        self.layer0, in_channels = self._create_layer(self.architecture["layer0"], in_channels)
        self.layer1, in_channels = self._create_layer(self.architecture["layer1"], in_channels)
        self.layer2, in_channels = self._create_layer(self.architecture["layer2"], in_channels)
        self.layer3, in_channels = self._create_layer(self.architecture["layer3"], in_channels)
        self.layer4, in_channels = self._create_layer(self.architecture["layer4"], in_channels)
        self.layer5, in_channels = self._create_layer(self.architecture["layer5"], in_channels)
        # fmt: on

        self.conv = nn.Conv2d(in_channels, num_classes, kernel_size=1, stride=1)
        self.avgpool = GlobalAvgPool2d()
        self.softmax = nn.Softmax(dim=1)

    def forward(self, x):
        x = self.layer0(x)
        x = self.layer1(x)
        x = self.layer2(x)
        x = self.layer3(x)
        x = self.layer4(x)
        x = self.layer5(x)

        x = self.conv(x)
        x = self.avgpool(x)
        x = self.softmax(x)

        return x

    def _create_layer(self, cfg, in_channels):
        modules = []
        for x in cfg:
            if type(x) == tuple:
                modules += create_cnn_with_bn_leaky(
                    in_channels, x[0], kernel_size=x[1], stride=x[2], padding=x[3]
                )
                in_channels = x[0]
            elif type(x) == str:
                modules += [nn.MaxPool2d(kernel_size=(2, 2), stride=(2, 2))]

        return nn.Sequential(*modules), in_channels
